{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "88216e84",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/openai_agent_retrieval.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "99cea58c-48bc-4af6-8358-df9695659983",
   "metadata": {},
   "source": [
    "# Retrieval-Augmented OpenAI Agent"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "673df1fe-eb6c-46ea-9a73-a96e7ae7942e",
   "metadata": {},
   "source": [
    "In this tutorial, we show you how to use our `FnRetrieverOpenAI` implementation\n",
    "to build an agent on top of OpenAI's function API and store/index an arbitrary number of tools. Our indexing/retrieval modules help to remove the complexity of having too many functions to fit in the prompt."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "54b7bc2e-606f-411a-9490-fcfab9236dfc",
   "metadata": {},
   "source": [
    "## Initial Setup "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23e80e5b-aaee-4f23-b338-7ae62b08141f",
   "metadata": {},
   "source": [
    "Let's start by importing some simple building blocks.  \n",
    "\n",
    "The main thing we need is:\n",
    "1. the OpenAI API\n",
    "2. a place to keep conversation history \n",
    "3. a definition for tools that our agent can use."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4ce34f4d",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "346fb488",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d47283b-025e-4874-88ed-76245b22f82e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/suo/miniconda3/envs/llama/lib/python3.9/site-packages/deeplake/util/check_latest_version.py:32: UserWarning: A newer version of deeplake (3.6.7) is available. It's recommended that you update to the latest version using `pip install -U deeplake`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "from typing import Sequence\n",
    "\n",
    "from llama_index.tools import BaseTool, FunctionTool"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6fe08eb1-e638-4c00-9103-5c305bfacccf",
   "metadata": {},
   "source": [
    "Let's define some very simple calculator tools for our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dd3c4a6-f3e0-46f9-ad3b-7ba57d1bc992",
   "metadata": {},
   "outputs": [],
   "source": [
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def useless(a: int, b: int) -> int:\n",
    "    \"\"\"Toy useless function.\"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply, name=\"multiply\")\n",
    "useless_tools = [\n",
    "    FunctionTool.from_defaults(fn=useless, name=f\"useless_{str(idx)}\")\n",
    "    for idx in range(28)\n",
    "]\n",
    "add_tool = FunctionTool.from_defaults(fn=add, name=\"add\")\n",
    "\n",
    "all_tools = [multiply_tool] + [add_tool] + useless_tools\n",
    "all_tools_map = {t.metadata.name: t for t in all_tools}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8170dd32-fa00-458c-aeb5-37292d0773c0",
   "metadata": {},
   "source": [
    "## Building an Object Index\n",
    "\n",
    "We have an `ObjectIndex` construct in LlamaIndex that allows the user to use our index data structures over arbitrary objects.\n",
    "The ObjectIndex will handle serialiation to/from the object, and use an underying index (e.g. VectorStoreIndex, SummaryIndex, KeywordTableIndex) as the storage mechanism. \n",
    "\n",
    "In this case, we have a large collection of Tool objects, and we'd want to define an ObjectIndex over these Tools.\n",
    "\n",
    "The index comes bundled with a retrieval mechanism, an `ObjectRetriever`. \n",
    "\n",
    "This can be passed in to our agent so that it can \n",
    "perform Tool retrieval during query-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6704a755-7f05-43a3-8a56-f5f587ae4c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an \"object\" index over these tools\n",
    "from llama_index import VectorStoreIndex\n",
    "from llama_index.objects import ObjectIndex, SimpleToolNodeMapping\n",
    "\n",
    "tool_mapping = SimpleToolNodeMapping.from_objects(all_tools)\n",
    "obj_index = ObjectIndex.from_objects(\n",
    "    all_tools,\n",
    "    tool_mapping,\n",
    "    VectorStoreIndex,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "707d30b8-6405-4187-a9ed-6146dcc42167",
   "metadata": {},
   "source": [
    "## Our `FnRetrieverOpenAIAgent` Implementation "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "798ca3fd-6711-4c0c-a853-d868dd14b484",
   "metadata": {},
   "source": [
    "We provide a `FnRetrieverOpenAIAgent` implementation in LlamaIndex, which can take in an `ObjectRetriever` over a set of `BaseTool` objects.\n",
    "\n",
    "During query-time, we would first use the `ObjectRetriever` to retrieve a set of relevant Tools. These tools would then be passed into the agent; more specifically, their function signatures would be passed into the OpenAI Function calling API. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ab3938-1138-43ea-b085-f430b42f5377",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import FnRetrieverOpenAIAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d852ece7-e5a1-4368-9d59-c7014e0b5b4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = FnRetrieverOpenAIAgent.from_retriever(\n",
    "    obj_index.as_retriever(), verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ea069f-819b-4ec1-a93c-fcbaacb362a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: multiply with args: {\n",
      "  \"a\": 212,\n",
      "  \"b\": 122\n",
      "}\n",
      "Got output: 25864\n",
      "========================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Response(response='212 multiplied by 122 is 25,864.', source_nodes=[], metadata=None)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.chat(\"What's 212 multiplied by 122? Make sure to use Tools\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec423b90-59cd-40ef-b497-a3842b3e7b58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: add with args: {\n",
      "  \"a\": 212,\n",
      "  \"b\": 122\n",
      "}\n",
      "Got output: 334\n",
      "========================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Response(response='212 added to 122 is 334.', source_nodes=[], metadata=None)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.chat(\"What's 212 added to 122 ? Make sure to use Tools\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
